import pandas as pd
import numpy as np
import panel as pn
import hvplot.pandas
import matplotlib.pyplot as plt
import holoviews as hv
from holoviews import opts
from scipy import optimize
import warnings
warnings.filterwarnings('ignore')

evidences = pd.read_csv('../output/evidence_aligned.csv')
evidences['lastAA'] = evidences['Sequence'].str[-1:]
evidences_trp = evidences[evidences.lastAA.isin(['K', 'R'])]

# Calculate trend line functions

CCS_fit_charge2 = evidences[evidences['Charge'] == 2] 
CCS_fit_charge3 = evidences[evidences['Charge'] == 3] 
CCS_fit_charge4 = evidences[evidences['Charge'] == 4]

def trendline_func(x, a, b):
    return a * np.power(x, b)

params_charge2, params_covariance_charge2 = optimize.curve_fit(
    trendline_func, CCS_fit_charge2['m/z'], CCS_fit_charge2['CCS'])
params_charge3, params_covariance_charge3 = optimize.curve_fit(
    trendline_func, CCS_fit_charge3['m/z'], CCS_fit_charge3['CCS'])
params_charge4, params_covariance_charge4 = optimize.curve_fit(
    trendline_func, CCS_fit_charge4['m/z'], CCS_fit_charge4['CCS'])

aminoacids = 'A R N D C Q E G H I L K M F P S T W Y V'.split()

# Figure 1
y_label = 'CCS (\u212b\u00B2)'
select_aa = pn.widgets.Select(name='Select an amino acid', options=aminoacids)

@pn.depends(select_aa.param.value)
def calculate_aa_content(aa):
    
    evidences_trp['count'] = evidences_trp['Sequence'].str.count(aa) * 100
    scatter_plot = evidences_trp.hvplot.points(x='m/z', y='CCS', c='count', 
                                               xlabel='m/z', ylabel=y_label, clabel= aa + '(%)', 
                                               width=450, height=400, cmap=plt.get_cmap('YlOrRd_r'), tools=['hover'],
                                               colorbar=True, rasterize=True)
    scatter_plot.opts(
        toolbar='above', clim=(0, 100)
    )
    
    opts.defaults(opts.Curve(color = "black", line_width=0.5, line_dash='dashed'))
    trendline_x = np.arange(300,1800,1)
    trendline_charge2 = hv.Curve((trendline_x, trendline_func(trendline_x, params_charge2[0], params_charge2[1])))
    trendline_charge3 = hv.Curve((trendline_x, trendline_func(trendline_x, params_charge3[0], params_charge3[1])))
    trendline_charge4 = hv.Curve((trendline_x, trendline_func(trendline_x, params_charge4[0], params_charge4[1])))
    
    return (scatter_plot * trendline_charge2 * trendline_charge3 * trendline_charge4)

layout = pn.Column(
    pn.pane.Markdown("""### Figure 3. A global view on peptide cross sections. b, Subset of peptides with C-terminal arginine or lysine colored by the fraction of amino acids in the linear sequence(n = 452,592).""", 
                     margin=(10, 0, 20, 0),
                     css_classes=['title'],),
    pn.Row(
        pn.layout.VSpacer(width=110),
        select_aa, 
        calculate_aa_content
    ),
    align='center',
    width=970
)
layout.embed()